import re
import json
import os
import re
from typing import List, Dict, Optional, Tuple
from langchain_core.documents import Document
from datetime import datetime
from ..core.logger import logger


class MetadataEnhancer:
    """文档元数据提取和增强器"""
    
    def __init__(self):
        """初始化元数据增强器"""
        # 公司名称模式匹配
        self.company_patterns = [
            r'([\u4e00-\u9fa5]{2,20}?(?:股份有限公司|有限公司|集团有限公司|集团|科技有限公司|银行|保险|证券))',
            r'([\u4e00-\u9fa5]{2,20}?(?:公司|集团|银行|保险|证券))',
        ]
        
        # 招股书发布日期模式
        self.date_patterns = [
            r'(\d{4}年\d{1,2}月\d{1,2}日)',
            r'(\d{4}-\d{2}-\d{2})',
            r'(\d{8})',  # 20230101格式
        ]
        
        # 章节层级模式
        self.chapter_level_patterns = {
            'chapter': r'第[一二三四五六七八九十百千]+章',  # 第X章
            'section': r'第[一二三四五六七八九十百千]+节',  # 第X节
            'subsection': r'[一二三四五六七八九十百千]+、',  # X、
            'item': r'\d+\.',  # 1.
        }
        
        # 内容类型关键词
        self.content_type_keywords = {
            '术语定义': ['指', '定义', '简称', '全称'],
            '业务描述': ['主营业务', '业务范围', '经营范围', '主要产品'],
            '财务数据': ['营业收入', '净利润', '资产', '负债', '现金流', '财务报表'],
            '风险提示': ['风险', '不确定性', '提示'],
            '募集资金': ['募集资金', '资金运用', '投资项目'],
            '股权结构': ['股权', '股东', '持股比例', '股权结构'],
            '管理层': ['董事', '监事', '高级管理人员', '管理层'],
        }
        
        # 页面编号模式
        self.page_number_patterns = [
            r'页码：?(\d+)',
            r'第\s*(\d+)\s*页',
        ]
        
        # 公司名称映射表（用于标准化）
        self.company_name_mapping = self._load_company_mapping()
    
    def _load_company_mapping(self) -> Dict[str, str]:
        """加载公司名称映射表"""
        # 使用绝对路径确保正确加载配置文件
        mapping_file = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), 'config', 'company_aliases.json')
        mapping = {}
        
        try:
            # 尝试从文件加载公司名称映射表
            if os.path.exists(mapping_file):
                with open(mapping_file, 'r', encoding='utf-8') as f:
                    aliases_data = json.load(f)
                
                # 支持两种格式：
                # 格式1: {"Standard Name": ["Alias1", "Alias2"]}
                # 格式2: [{"standard_name": "Standard Name", "aliases": ["Alias1", "Alias2"]}]
                
                if isinstance(aliases_data, dict):
                    # 处理格式1
                    for standard_name, aliases in aliases_data.items():
                        if aliases and isinstance(aliases, list):
                            for alias in aliases:
                                mapping[alias] = standard_name
                                
                elif isinstance(aliases_data, list):
                    # 处理格式2
                    for item in aliases_data:
                        if isinstance(item, dict) and 'standard_name' in item and 'aliases' in item:
                            standard_name = item['standard_name']
                            aliases = item['aliases']
                            if aliases and isinstance(aliases, list):
                                for alias in aliases:
                                    mapping[alias] = standard_name
                
                logger.info(f"成功从 {mapping_file} 加载公司名称映射表，共 {len(mapping)} 个映射关系")
            else:
                logger.warning(f"未找到公司名称映射文件: {mapping_file}")
        except Exception as e:
            logger.error(f"加载公司名称映射表时出错: {e}")
            
        return mapping
    
    def extract_company_name(self, text: str) -> Optional[str]:
        """从文本中提取公司名称"""
        for pattern in self.company_patterns:
            matches = re.findall(pattern, text)
            if matches:
                # 选择出现频率最高或最长的公司名称
                company_name = max(matches, key=lambda x: len(x))
                # 应用公司名称映射表进行标准化
                return self.company_name_mapping.get(company_name, company_name)
        return None
    
    def extract_date(self, text: str) -> Optional[str]:
        """从文本中提取日期"""
        for pattern in self.date_patterns:
            match = re.search(pattern, text)
            if match:
                date_str = match.group(1)
                # 尝试解析日期格式
                try:
                    # 处理不同的日期格式
                    if '-' in date_str:
                        date_obj = datetime.strptime(date_str, '%Y-%m-%d')
                    elif '年' in date_str:
                        # 处理中文日期格式
                        match_parts = re.search(r'(\d{4})年(\d{1,2})月(\d{1,2})日', date_str)
                        if match_parts:
                            year, month, day = map(int, match_parts.groups())
                            date_obj = datetime(year, month, day)
                        else:
                            continue
                    elif len(date_str) == 8:
                        date_obj = datetime.strptime(date_str, '%Y%m%d')
                    else:
                        continue
                    # 返回标准化的日期字符串
                    return date_obj.strftime('%Y-%m-%d')
                except ValueError:
                    continue
        return None
    
    def extract_chapter_hierarchy(self, text: str) -> Dict[str, str]:
        """提取章节层级信息"""
        hierarchy = {}
        
        for level, pattern in self.chapter_level_patterns.items():
            match = re.search(pattern, text[:200])
            if match:
                hierarchy[level] = match.group(0)
        
        return hierarchy
    
    def classify_content_type(self, text: str) -> List[str]:
        """对文本内容类型进行分类"""
        content_types = []
        
        for content_type, keywords in self.content_type_keywords.items():
            for keyword in keywords:
                if keyword in text:
                    content_types.append(content_type)
                    break  # 每个内容类型只需要匹配一个关键词
        
        # 如果没有匹配到任何类型，标记为其他
        if not content_types:
            content_types.append('其他')
            
        return content_types
    
    def extract_page_number(self, text: str) -> Optional[str]:
        """提取页面编号"""
        for pattern in self.page_number_patterns:
            match = re.search(pattern, text)
            if match:
                return match.group(1)
        return None
    
    def enhance_metadata(self, document: Document) -> Document:
        """增强单个文档的元数据"""
        text = document.page_content
        metadata = document.metadata.copy()
        
        # 提取并添加公司名称
        if 'company_name' not in metadata or not metadata['company_name']:
            company_name = self.extract_company_name(text)
            if company_name:
                metadata['company_name'] = company_name
        
        # 提取并添加发布日期
        if 'publish_date' not in metadata or not metadata['publish_date']:
            publish_date = self.extract_date(text)
            if publish_date:
                metadata['publish_date'] = publish_date
        
        # 提取并添加章节层级
        chapter_hierarchy = self.extract_chapter_hierarchy(text)
        if chapter_hierarchy:
            metadata.update(chapter_hierarchy)
        
        # 提取并添加内容类型
        content_types = self.classify_content_type(text)
        if content_types:
            metadata['content_type'] = content_types
        
        # 提取并添加页面编号
        if 'page_number' not in metadata or not metadata['page_number']:
            page_number = self.extract_page_number(text)
            if page_number:
                metadata['page_number'] = page_number
        
        # 添加文档来源信息
        if 'source' in metadata:
            # 从文件路径中提取更多信息
            source_path = metadata['source']
            file_name = os.path.basename(source_path)
            metadata['file_name'] = file_name
            
            # 尝试从文件名中提取更多信息
            file_company = self.extract_company_name(file_name)
            if file_company and 'company_name' not in metadata:
                metadata['company_name'] = file_company
        
        # 添加增强标记
        metadata['metadata_enhanced'] = True
        
        # 创建新的文档对象
        enhanced_doc = Document(
            page_content=text,
            metadata=metadata
        )
        
        return enhanced_doc
    
    def batch_enhance_metadata(self, documents: List[Document]) -> List[Document]:
        """批量增强文档的元数据"""
        enhanced_docs = []
        
        for doc in documents:
            enhanced_doc = self.enhance_metadata(doc)
            enhanced_docs.append(enhanced_doc)
        
        return enhanced_docs

    def save_enhanced_docs(self, docs: List[Document], output_dir: str):
        """保存增强后的文档到文件"""
        os.makedirs(output_dir, exist_ok=True)
        
        for i, doc in enumerate(docs):
            file_name = f"doc_{i}.json"
            file_path = os.path.join(output_dir, file_name)
            
            with open(file_path, 'w', encoding='utf-8') as f:
                json.dump({
                    'page_content': doc.page_content,
                    'metadata': doc.metadata
                }, f, ensure_ascii=False, indent=2)


def extract_title_from_content(content: str) -> str:
    """
    从文档内容中提取标题
    :param content: 文档内容字符串
    :return: 提取的标题，如果没有找到合适的标题则返回空字符串
    """
    if not content:
        return ""
    
    lines = content.split('\n')
    
    # 遍历前10行寻找可能的标题
    for i, line in enumerate(lines[:10]):
        line = line.strip()
        # 标题特征：长度适中（10-200字符），可能是第一行，可能包含章节标记
        if (line and 
            10 <= len(line) <= 200 and 
            not line.startswith('=') and 
            not line.startswith('-') and 
            not line.startswith(' ') and 
            not line.isupper()):  # 避免全大写的页眉页脚
            # 检查是否包含章节标记
            if re.match(r'^第[一二三四五六七八九十百千]+[章节篇]', line) or \
               re.match(r'^[0-9]+\.[0-9]+', line) or \
               re.match(r'^[A-Za-z]\.', line):
                return line
            
            # 如果是第一行且不包含特殊标记，也可能是标题
            if i == 0 and len(line) > 15:
                return line
    
    return ""


def extract_keywords_from_content(content: str, max_keywords: int = 5) -> list:
    """
    从文档内容中提取关键词（简单实现）
    :param content: 文档内容字符串
    :param max_keywords: 最大关键词数量
    :return: 关键词列表
    """
    if not content:
        return []
    
    # 简单的关键词提取策略：
    # 1. 统计出现频率较高的名词
    # 2. 去除停用词
    
    # 基本的中文停用词列表
    stop_words = {
        '的', '了', '在', '是', '我', '有', '和', '就', '不', '人', '都', '一', '一个', '上', '也', 
        '很', '到', '说', '要', '去', '你', '会', '着', '没有', '看', '好', '自己', '这', 
        '对于', '关于', '因此', '所以', '但是', '而且', '因为', '例如', '包括', '特别是',
        '主要', '之一', '首先', '其次', '最后', '总之', '综上所述'
    }
    
    # 分割文本并统计词频
    # 这里使用简单的正则表达式分割，实际应用中可以使用更复杂的分词工具
    words = re.findall(r'[\u4e00-\u9fa5]{2,}|[A-Za-z]+', content)
    word_freq = {}
    
    for word in words:
        # 过滤掉停用词和太短的词
        if word not in stop_words and len(word) >= 2:
            word_freq[word] = word_freq.get(word, 0) + 1
    
    # 按词频排序并取前N个
    sorted_keywords = sorted(word_freq.items(), key=lambda x: x[1], reverse=True)
    return [word for word, freq in sorted_keywords[:max_keywords]]


def extract_sections_from_content(content: str) -> list:
    """
    从文档内容中提取章节信息
    :param content: 文档内容字符串
    :return: 章节列表，每个元素包含章节标题和起始位置
    """
    if not content:
        return []
    
    sections = []
    lines = content.split('\n')
    
    # 匹配常见的章节格式
    section_patterns = [
        r'^第[一二三四五六七八九十百千]+[章节篇]\s*[:：]?\s*(.+)',  # 中文数字章节
        r'^[0-9]+\.[0-9]+\s+(.+)',  # 1.1 格式
        r'^[0-9]+\s+(.+)',  # 1 格式
        r'^[A-Za-z]\)\s+(.+)',  # A) 格式
        r'^[一二三四五六七八九十]+\)\s+(.+)',  # 一) 格式
    ]
    
    for i, line in enumerate(lines):
        line = line.strip()
        if not line or len(line) > 200:  # 章节标题通常不会太长
            continue
        
        # 检查是否匹配章节格式
        for pattern in section_patterns:
            match = re.match(pattern, line)
            if match:
                # 计算章节在原始内容中的大致位置
                position = sum(len(l) + 1 for l in lines[:i])  # +1 是换行符
                sections.append({
                    'title': line,
                    'position': position
                })
                break
    
    return sections


def enhance_document_metadata(doc: Document) -> Document:
    """
    增强文档的元数据
    :param doc: 原始文档对象
    :return: 增强后的文档对象
    """
    # 提取标题
    title = extract_title_from_content(doc.page_content)
    if title:
        doc.metadata['title'] = title
        logger.debug(f"提取到标题: {title}")
    
    # 提取关键词（如果内容足够长）
    if len(doc.page_content) > 300:  # 只有当内容足够长时才提取关键词
        keywords = extract_keywords_from_content(doc.page_content)
        if keywords:
            doc.metadata['keywords'] = keywords
            logger.debug(f"提取到关键词: {keywords}")
    
    # 提取章节信息（如果内容足够长）
    if len(doc.page_content) > 1000:  # 只有当内容足够长时才提取章节
        sections = extract_sections_from_content(doc.page_content)
        if sections:
            # 只保存章节标题列表，不保存位置信息
            doc.metadata['sections'] = [s['title'] for s in sections]
            logger.debug(f"提取到章节: {len(sections)} 个")
    
    # 补充其他有用的元数据
    # 1. 内容长度
    doc.metadata['content_length'] = len(doc.page_content)
    
    # 2. 估算页数（假设每页约1000字符）
    estimated_pages = max(1, len(doc.page_content) // 1000)
    doc.metadata['estimated_pages'] = estimated_pages
    
    # 3. 检查是否包含表格数据的标记
    if any(keyword in doc.page_content.lower() for keyword in ['表格', 'table', '列表', 'list', '数据']):
        doc.metadata['contains_tables'] = True
    
    return doc